package cn.tedu.util;


import java.util.concurrent.*;
import java.util.function.Function;
import java.util.function.Supplier;

public class AsyncTimeoutUtil {
    private static final ScheduledExecutorService scheduler =
            Executors.newScheduledThreadPool(1, r -> {
                Thread t = new Thread(r);
                t.setDaemon(true);
                return t;
            });

    public static <T> CompletableFuture<T> withTimeout(
            Supplier<T> supplier,
            long timeout,
            TimeUnit unit,
            Executor executor) {

        CompletableFuture<T> taskFuture = CompletableFuture.supplyAsync(supplier, executor);
        CompletableFuture<T> timeoutFuture = new CompletableFuture<>();

        scheduler.schedule(() ->
                        timeoutFuture.completeExceptionally(
                                new TimeoutException("Operation timed out after " + timeout + " " + unit)),
                timeout, unit);

        return taskFuture.applyToEither(timeoutFuture, Function.identity())
                .exceptionally(ex -> {
                    if (ex instanceof CompletionException &&
                            ex.getCause() instanceof TimeoutException) {
                        taskFuture.cancel(true);
                    }
                    throw ex instanceof CompletionException ?
                            (CompletionException) ex : new CompletionException(ex);
                });
    }
}
